#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import serial
import rclpy
import math
import serial.tools.list_ports
from sensor_msgs.msg import Imu
from sensor_msgs.msg import MagneticField
import time
import modbus_tk
import modbus_tk.defines as cst
from modbus_tk import modbus_rtu

from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import Quaternion
import tf_transformations

# 查找 ttyUSB* 设备
def find_ttyUSB():
    print('imu 默认串口为 /dev/ttyUSB0, 若识别多个串口设备, 请在 launch 文件中修改 imu 对应的串口')
    posts = [port.device for port in serial.tools.list_ports.comports() if 'USB' in port.device]
    print('当前电脑所连接的 {} 串口设备共 {} 个: {}'.format('USB', len(posts), posts))


angularVelocity = [0, 0, 0]
acceleration = [0, 0, 0]
magnetometer = [0, 0, 0]
angle_degree = [0, 0, 0]


def main(args=None):
    rclpy.init(args=args)

    find_ttyUSB()

    node = rclpy.create_node("imu")
    
    #使用declare_parameter方法来声明你想要使用的参数
    node.declare_parameter('port', '')
    node.declare_parameter('baud',0)

    #使用get_parameter方法来获取参数的值
    port = node.get_parameter('port').get_parameter_value().string_value
    baudrate = node.get_parameter('baud').get_parameter_value().integer_value

    print("IMU Type: Modbus Port:%s baud:%d" % (port, baudrate))
    imu_pub = node.create_publisher(Imu, "wit/imu", 10)
    mag_pub = node.create_publisher(MagneticField, "wit/mag", 10)

    try:
        wt_imu = serial.Serial(port=port, baudrate=baudrate, timeout=0.5)
        if wt_imu.isOpen():
            node.get_logger().info("\033[32mport open success...\033[0m")
        else:
            wt_imu.open()
            node.get_logger().info("\033[32mport open success...\033[0m")
    except Exception as e:
        print(e)
        node.get_logger().info("\033[31mport open failed\033[0m")
        exit(0)

    master = modbus_rtu.RtuMaster(wt_imu)
    master.set_timeout(1)
    master.set_verbose(True)

    imu_msg = Imu()
    mag_msg = MagneticField()

    while rclpy.ok():
        try:
            reg = master.execute(80, cst.READ_HOLDING_REGISTERS, 52, 12)
        except Exception as e:
            print(e)
            node.get_logger().info("\033[31mread register time out, please check connection or baundrate set!\033[0m")
            time.sleep(0.1)
        else:
            v = [0] * 12
            for i in range(0, 12):
                if (reg[i] > 32767):
                    v[i] = reg[i] - 65536
                else:
                    v[i] = reg[i]

            acceleration = [v[i] / 32768.0 * 16 * 9.8 for i in range(0, 3)]
            angularVelocity = [v[i] / 32768.0 * 2000 * math.pi / 180 for i in range(3, 6)]
            magnetometer = v[6:9]
            angle_degree = [v[i] / 32768.0 * 180 for i in range(9, 12)]

            angle_radian = [angle_degree[i] * math.pi / 180 for i in range(3)]
            qua = tf_transformations.quaternion_from_euler(angle_radian[0], angle_radian[1], angle_radian[2])

            # 获取当前时间戳
            stamp = node.get_clock().now().to_msg()

            # 设置时间戳和帧ID
            imu_msg.header.stamp = stamp
            imu_msg.header.frame_id = "base_link"

            mag_msg.header.stamp = stamp
            mag_msg.header.frame_id = "base_link"

            imu_msg.orientation = Quaternion(x=qua[0], y=qua[1], z=qua[2], w=qua[3])

            imu_msg.angular_velocity.x = angularVelocity[0]
            imu_msg.angular_velocity.y = angularVelocity[1]
            imu_msg.angular_velocity.z = angularVelocity[2]

            imu_msg.linear_acceleration.x = acceleration[0]
            imu_msg.linear_acceleration.y = acceleration[1]
            imu_msg.linear_acceleration.z = acceleration[2]

            mag_msg.magnetic_field.x = float(magnetometer[0])
            mag_msg.magnetic_field.y = float(magnetometer[1])
            mag_msg.magnetic_field.z = float(magnetometer[2])

            imu_pub.publish(imu_msg)
            mag_pub.publish(mag_msg)

if __name__ == '__main__':
    main()




